#pragma once

#include "gtsam/base/Manifold.h"
#include "gtsam/base/Vector.h"
#include "gtsam/geometry/Pose3.h"

/**
 * NavState: attitude, position and velocity
 */

namespace gtsam {
/**
 * Navigation state: Pose (Rotation, translation) and velocity
 * NOTE(frank): it does not make sense to make this a Lie Group, but it is a 9D
 * manifold
 */
class NavState {
 private:
  Rot3 R_;     // Rotation nRb, rotates points in body frame to nav frame
  Point3 t_;   // position in nav frame
  Vector3 v_;  // velocity in nav frame

 public:
  enum { dimension = 9 };

  typedef std::pair<Point3, Vector3> PositionAndVelocity;

  /**
   * constructors
   */
  // default
  NavState() : t_(0, 0, 0), v_(0, 0, 0) {}
  // construct from attitude, position and velocity
  NavState(const Rot3& R, const Point3& t, const Vector3& v)
      : R_(R), t_(t), v_(v) {}
  // construct from pose and velocity
  NavState(const Pose3& pose, const Vector3& v)
      : R_(pose.rotation()), t_(pose.translation()), v_(v) {}
  /// construct from SO(3) and R^6
  NavState(const Matrix3& R, const Vector6& tv)
      : R_(R), t_(tv.head<3>()), v_(tv.tail<3>()) {}
  // constructor with derivatives
  static NavState Create(const Rot3& R,
                         const Point3& t,
                         const Vector3& v,
                         OptionalJacobian<9, 3> HR,
                         OptionalJacobian<9, 3> HT,
                         OptionalJacobian<9, 3> HV);
  static NavState FromPoseVelocity(const Pose3& pose,
                                   const Vector3& v,
                                   OptionalJacobian<9, 6> HP,
                                   OptionalJacobian<9, 3> HV);

  /**
   * component access
   */
  const Rot3& attitude(OptionalJacobian<3, 9> H = boost::none) const;
  const Point3& position(OptionalJacobian<3, 9> H = boost::none) const;
  const Vector3& velocity(OptionalJacobian<3, 9> H = boost::none) const;

  const Pose3 pose() const {
    return Pose3(attitude(), position());
  }

  /**
   * derived quantities
   */
  Matrix3 R() const {
    return R_.matrix();
  }
  // return quaternion. Induces computation in matrix mode
  Quaternion quaternion() const {
    return R_.toQuaternion();
  }
  // return position as Vector3
  Vector3 t() const {
    return t_;
  }
  // return velocity as Vector3. Computation-free.
  const Vector3& v() const {
    return v_;
  }
  // return velocity in body frame
  Vector3 bodyVelocity(OptionalJacobian<3, 9> H = boost::none) const;

  /// Return matrix group representation, in MATLAB notation:
  /// nTb = [nRb 0 n_t; 0 nRb n_v; 0 0 1]
  /// With this embedding in GL(3), matrix product agrees with compose
  Matrix7 matrix() const;

  /**
   * Testable
   */

  /// Output stream operator
  GTSAM_EXPORT
  friend std::ostream& operator<<(std::ostream& os, const NavState& state);

  /// print
  void print(const std::string& s = "") const;

  /// equals
  bool equals(const NavState& other, double tol = 1e-8) const;

  /**
   * Manifold
   */
  // Tangent space sugar.
  // TODO(frank): move to private navstate namespace in cpp
  static Eigen::Block<Vector9, 3, 1> dR(Vector9& v) {
    return v.segment<3>(0);
  }
  static Eigen::Block<Vector9, 3, 1> dP(Vector9& v) {
    return v.segment<3>(3);
  }
  static Eigen::Block<Vector9, 3, 1> dV(Vector9& v) {
    return v.segment<3>(6);
  }
  static Eigen::Block<const Vector9, 3, 1> dR(const Vector9& v) {
    return v.segment<3>(0);
  }
  static Eigen::Block<const Vector9, 3, 1> dP(const Vector9& v) {
    return v.segment<3>(3);
  }
  static Eigen::Block<const Vector9, 3, 1> dV(const Vector9& v) {
    return v.segment<3>(6);
  }

  // retract with optional derivatives
  NavState retract(const Vector9& v,
                   OptionalJacobian<9, 9> H1 = boost::none,
                   OptionalJacobian<9, 9> H2 = boost::none) const;

  // localCoordinates with optional derivatives
  Vector9 localCoordinates(const NavState& g,
                           OptionalJacobian<9, 9> H1 = boost::none,
                           OptionalJacobian<9, 9> H2 = boost::none) const;

  /**
   * Dynamics
   */

  // integrate forward in time given angular velocity and acceleration in body
  // frame uses second order integration for position, return derivatives except
  // dt.
  NavState update(const Vector3& b_acceleration,
                  const Vector3& b_omega,
                  const double dt,
                  OptionalJacobian<9, 9> F,
                  OptionalJacobian<9, 3> G1,
                  OptionalJacobian<9, 3> G2) const;

  // compute tangent space contribution due to Coriolis forces
  Vector9 coriolis(double dt,
                   const Vector3& omega,
                   bool secondOrder = false,
                   OptionalJacobian<9, 9> H = boost::none) const;

  // correct preintegrated tangent vector with our velocity and rotated gravity.
  // taking into account Coriolis forces if omegaCoriolis is given
  Vector9 correctPIM(const Vector9& pim,
                     double dt,
                     const Vector3& n_gravity,
                     const boost::optional<Vector3>& omegaCoriolis,
                     bool use2OrderCoriolis = false,
                     OptionalJacobian<9, 9> H1 = boost::none,
                     OptionalJacobian<9, 9> H2 = boost::none) const;
};

// specialize NavState traits to use a Retract/Local that agrees with IMUFactors
template <>
struct traits<NavState> : internal::Manifold<NavState> {};

}  // namespace gtsam
